# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import logging
import signal
import string
import random
import os

from coin.base.datetime_util import parse_end_time
from coin.flow.simple_queue import SimpleQueue
from tornado.ioloop import IOLoop


def run(subscriber_init,
        queue_name=None,
        archive_root_dir=None,
        ioloop=None,
        install_signal_handler=True):
  begin_time = datetime.datetime.now()
  end_time = parse_end_time('00:00:05')
  print('End time: %s' % end_time)

  if queue_name is None:
    random_str = \
        ''.join(random.choices(string.ascii_lowercase, k=1) +
                random.choices(string.ascii_lowercase + string.digits, k=5))
    queue_name = 'queue-%s%s' % (begin_time.strftime('%H%M'), random_str)
  queue = SimpleQueue(queue_name)

  archive_root_dir = archive_root_dir or '~/data/flow'
  archiver = queue.start_archiver(queue_time=begin_time, root_dir=archive_root_dir)

  ioloop = ioloop or IOLoop.current()

  def sig_handler(self, sig=None, frame=None):
    ioloop.add_callback_from_signal(ioloop.stop)

  subscriber = subscriber_init(queue)
  try:
    subscriber.start(ioloop)
    ioloop.add_timeout(end_time - datetime.datetime.now(), ioloop.stop)
    if install_signal_handler:
      ioloop.add_callback(signal.signal, signal.SIGINT, sig_handler)
    ioloop.start()
  finally:
    subscriber.stop()
    archiver.stop()
  return 0


def run2(init,
         teardown,
         queue_name=None,
         archive_root_dir=None,
         ioloop=None,
         install_signal_handler=True,
         end_time=None,
         logger=None):
  begin_time = datetime.datetime.now()
  logger = logger or logging.getLogger(__name__)
  if end_time is not None:
    logger.info('End time: %s' % end_time)

  if queue_name is None:
    random_str = \
        ''.join(random.choices(string.ascii_lowercase, k=1) +
                random.choices(string.ascii_lowercase + string.digits, k=5))
    queue_name = 'queue-%s%s' % (begin_time.strftime('%H%M'), random_str)
  queue = SimpleQueue(queue_name)

  archive_root_dir = archive_root_dir or '~/data/flow'
  archiver = queue.start_archiver(queue_time=begin_time, root_dir=archive_root_dir)

  ioloop = ioloop or IOLoop.current()

  def sig_handler(self, sig=None, frame=None):
    ioloop.add_callback_from_signal(ioloop.stop)

  try:
    init(queue, ioloop)
    if end_time is not None:
      ioloop.add_timeout(end_time - datetime.datetime.now(), ioloop.stop)
    if install_signal_handler:
      ioloop.add_callback(signal.signal, signal.SIGINT, sig_handler)
    ioloop.start()
  finally:
    try:
      teardown()
    finally:
      archiver.stop()
      # For Kraken, ThreadPoolExecutor is used and up to 20 threads are created
      # to query feed. Those threads will take long time to finish after stop.
      # Here forcely quit instead of pending.
      os._exit(0)
